from enum import Enum, unique
import numpy as np

from src.utils.transformation import Transformation


@unique
class MotionState(Enum):
    STAND = 0
    WALK = 1
    WALK2STAND = 2


@unique
class LegState(Enum):
    LST = 0
    RST = 1
    DST = 2


class RobotState:
    def __init__(self, model_nv: int, fe_rot=np.eye(3)) -> None:
        super().__init__()

        self.fe_L_rot_L_off = np.eye(3)
        self.fe_R_rot_L_off = np.eye(3)

        self.js_eul_des = np.zeros(3)
        self.js_pos_des = np.zeros(3)
        self.js_omega_des = np.zeros(3)
        self.js_vel_des = np.zeros(3)

        self.model_nv = model_nv
        self.dyn_M = np.zeros((model_nv, model_nv))
        self.dyn_M_inv = np.zeros((model_nv, model_nv))
        self.dyn_C = np.zeros((model_nv, model_nv))
        self.dyn_Ag = np.zeros((6, model_nv))
        self.dyn_dAg = np.zeros((6, model_nv))
        self.dyn_G = np.zeros(model_nv)
        self.dyn_Non = np.zeros(model_nv)

        self.base_omega_L = np.zeros(3)
        self.base_omega_W = np.zeros(3)
        self.base_rpy = np.zeros(3)

        self.slop = np.zeros(3)
        self.inertia = np.zeros((3, 3))

        self.motors_pos_des = np.zeros(self.model_nv - 6)
        self.motors_vel_des = np.zeros(self.model_nv - 6)
        self.motors_tor_des = np.zeros(self.model_nv - 6)
        self.motors_tor_out = np.zeros(self.model_nv - 6)

        self.q = np.zeros(self.model_nv + 1)
        self.dq = np.zeros(self.model_nv)
        self.ddq = np.zeros(self.model_nv)
        self.qOld = np.zeros(self.model_nv + 1)
        self.J_base = np.zeros((6, self.model_nv))
        self.J_l = np.zeros((6, self.model_nv))
        self.J_r = np.zeros((6, self.model_nv))
        self.J_hip_link = np.zeros((6, self.model_nv))
        self.dJ_base = np.zeros((6, self.model_nv))
        self.dJ_l = np.zeros((6, self.model_nv))
        self.dJ_r = np.zeros((6, self.model_nv))
        self.Jcom_W = np.zeros((3, self.model_nv))
        self.pCoM_W = np.zeros(3)
        self.fe_l_pos_W = np.zeros(3)
        self.fe_r_pos_W = np.zeros(3)
        self.base_pos = np.zeros(3)
        self.fe_l_rot_W = np.zeros((3, 3))
        self.fe_r_rot_W = np.zeros((3, 3))
        self.base_rot = np.zeros((3, 3))
        self.fe_r_pos_L = np.zeros(3)
        self.fe_l_pos_L = np.zeros(3)
        self.hip_link_pos = np.zeros(3)
        self.hip_l_pos_L = np.zeros(3)
        self.hip_r_pos_L = np.zeros(3)
        self.hip_l_pos_W = np.zeros(3)
        self.hip_r_pos_W = np.zeros(3)
        self.fe_l_rot_L = np.zeros((3, 3))
        self.fe_r_rot_L = np.zeros((3, 3))
        self.hip_link_rot = np.zeros((3, 3))

        self.qCmd = np.zeros(self.model_nv + 1)
        self.dqCmd = np.zeros(self.model_nv)
        self.tauJointCmd = np.zeros(self.model_nv - 6)

        self.motionState = MotionState.STAND
        self.inertia = np.zeros((3, 3))

        self.tSwing = 0.0
        self.swingStartPos_W = np.zeros(3)
        self.swingDesPosCur_W = np.zeros(3)
        self.swingDesPosCur_L = np.zeros(3)
        self.swingDesPosFinal_W = np.zeros(3)
        self.stanceDesPos_W = np.zeros(3)
        self.posHip_W = np.zeros(3)
        self.posST_W = np.zeros(3)
        self.theta0 = 0.0
        self.width_hips = 0.0
        self.legState = LegState.LST
        self.legStateNext = LegState.LST
        self.phi = 0.0
        self.FL_est = np.zeros(6)
        self.FR_est = np.zeros(6)

        self.rpy = np.zeros(3)
        self.fL = np.zeros(3)
        self.fR = np.zeros(3)
        self.basePos = np.zeros(3)
        self.baseLinVel = np.zeros(3)
        self.baseAcc = np.zeros(3)
        self.baseAngVel = np.zeros(3)
        self.motors_pos_cur = np.zeros(self.model_nv - 6)
        self.motors_vel_cur = np.zeros(self.model_nv - 6)
        self.motors_tor_cur = np.zeros(self.model_nv - 6)

        self.Xd = np.zeros(12 * 10)
        self.X_cur = np.zeros(12)
        self.X_cal = np.zeros(12)
        self.dX_cal = np.zeros(12)
        self.fe_react_tau_cmd = np.zeros(12)

        self.qp_nWSR_MPC = 0
        self.qp_cpuTime_MPC = 0.0
        self.qpStatus_MPC = 0

        self.base_rpy_des = np.zeros(3)
        self.base_pos_des = np.zeros(3)
        self.des_ddq = np.zeros(self.model_nv)
        self.des_dq = np.zeros(self.model_nv)
        self.des_delta_q = np.zeros(self.model_nv)
        self.des_q = np.zeros(self.model_nv)
        self.swing_fe_pos_des_W = np.zeros(3)
        self.swing_fe_rpy_des_W = np.zeros(3)
        self.stance_fe_pos_cur_W = np.zeros(3)
        self.stance_fe_rot_cur_W = np.zeros((3, 3))
        self.wbc_delta_q_final = np.zeros(self.model_nv)
        self.wbc_dq_final = np.zeros(self.model_nv)
        self.wbc_ddq_final = np.zeros(self.model_nv)
        self.wbc_tauJointRes = np.zeros(self.model_nv - 6)
        self.wbc_FrRes = np.zeros(12)
        self.Fr_ff = np.zeros(12)
        self.qp_nWSR = 0
        self.qp_cpuTime = 0.0
        self.qp_status = 0

        self.fe_rot = fe_rot.copy()
        self.foot_height = 0.032

    def update_state(self):
        self.base_omega_W[0: 3] = self.baseAngVel[0: 3]
        Rcur = Transformation.eul2Rot(self.rpy[0], self.rpy[1], self.rpy[2])
        self.base_omega_W[:] = Rcur @ self.base_omega_W

        quatNow = Transformation.eul2quat(self.rpy[0], self.rpy[1], self.rpy[2])
        self.q[0: 3] = self.basePos[0: 3]
        self.q[3: 6] = quatNow[1: 4]
        self.q[6] = quatNow[0]

        for i in range(self.model_nv - 6):
            self.q[i + 7] = self.motors_pos_cur[i]

        vCom_W = self.baseLinVel[0: 3]
        self.dq[0: 3] = vCom_W[:]
        self.dq[3: 6] = self.base_omega_W[0: 3]

        for i in range(self.model_nv - 6):
            self.dq[i + 6] = self.motors_vel_cur[i]

        self.base_pos[0: 3] = self.q[0: 3]
        self.base_rpy[0: 3] = self.rpy[0: 3]
        self.base_rot[:, :] = Rcur[:, :]
        self.qOld[:] = self.q[:]
